#!/usr/bin/env python3

# 整理：无涯有涯
# 用法：python3 <(curl -sSL https://gitee.com/veryx/galaxy-tools/raw/master/9-export/export-cluster-info.py)
import sys

if sys.version_info.major < 3:
    print("需要Python 3或更高版本")
    sys.exit(1)

import os
import re
import csv
import sys
import json
import time
from fnmatch import fnmatch
from datetime import datetime
from subprocess import run
from typing import Dict, List, Any, Tuple

# 配置常量
CONFIG = {
    "debug": False,
    "cnpg_cluster": {"namespace": "pg-system", "name": "kd-cosmic"},
    "container_filters": {
        "exclude_rules": [
            ("kurl", "*", "*"),
            ("defalut", "*", "*"),
            ("openebs", "*", "*"),
            ("monitoring", "*", "*"),
            ("cnpg-system", "*", "*"),
            ("flux-system", "*", "*"),
            ("kube-system", "*", "*"),
            ("metallb-system", "*", "*"),
            ("*", "*", "kps"),
            ("*", "*", "*metrics*"),
            ("*", "*", "sshserver"),
            ("*", "*", "*operator*"),
            ("*", "*", "config-reloader"),
            ("*", "*", "migration-server"),
            ("kd-cosmic-log", "*", "zookeeper"),
            ("kd-cosmic-ops", "*", "reflector"),
        ],
        "include_rules": [("zk-system", "*", "*zookeeper*")],
    },
    "cnpg_whitelist": [
        "shared_buffers",
        "effective_cache_size",
    ],
}


class Logger:
    """统一的日志记录器"""

    def __init__(self, debug_mode: bool = True):
        self.debug_mode = debug_mode

    def is_tty(self) -> bool:
        return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()

    def _log(self, level: str, color_code: str, message: str):
        prefix = f"\033[{color_code}m{level} \033[0m" if self.is_tty() else level
        print(f"{prefix}: {message}")

    def info(self, message: str):
        self._log("INF", "92", message)

    def warn(self, message: str):
        self._log("WAR", "93", message)

    def error(self, message: str):
        self._log("ERR", "91", message)

    def debug(self, message: str):
        if self.debug_mode:
            self._log("DBG", "94", message)


class KubectlClient:
    """Kubectl命令执行客户端"""

    def __init__(self, logger: Logger, timeout: int = 30, max_retries: int = 3):
        self.logger = logger
        self.timeout = timeout
        self.max_retries = max_retries

    def run_with_retry(self, cmd_args: List[str]) -> Dict[str, Any]:
        """带重试机制的kubectl命令执行"""
        for attempt in range(self.max_retries):
            try:
                result = run(cmd_args, capture_output=True, text=True, timeout=self.timeout)
                if result.returncode == 0:
                    return json.loads(result.stdout)
                else:
                    self.logger.warn(f"命令执行失败 (尝试 {attempt+1}/{self.max_retries}): {result.stderr}")
            except Exception as e:
                self.logger.error(f"执行异常 (尝试 {attempt+1}/{self.max_retries}): {e}")

            if attempt < self.max_retries - 1:
                time.sleep(2**attempt)  # 指数退避

        return {"items": []}

    def get_nodes(self) -> Dict[str, Any]:
        """获取节点信息"""
        return self.run_with_retry(["kubectl", "get", "nodes", "-o", "json"])

    def get_deployments(self) -> Dict[str, Any]:
        """获取所有Deployment"""
        return self.run_with_retry(["kubectl", "get", "deployments", "-A", "-o", "json"])

    def get_statefulsets(self) -> Dict[str, Any]:
        """获取所有StatefulSet"""
        return self.run_with_retry(["kubectl", "get", "statefulsets", "-A", "-o", "json"])

    def get_cnpg_cluster(self, namespace: str, name: str) -> Dict[str, Any]:
        """获取CNPG集群信息"""
        return self.run_with_retry(["kubectl", "get", "cluster", "-n", namespace, name, "-o", "json"])


class ResourceFormatter:
    """资源格式化工具类"""

    @staticmethod
    def convert_to_gi(num_with_units: str) -> str:
        """转换内存单位到Gi"""
        if not num_with_units:
            return ""

        units = {
            "Ki": 1 / (1024 * 1024),
            "Mi": 1 / 1024,
            "Gi": 1,
        }

        for unit, factor in units.items():
            if unit in num_with_units:
                try:
                    value = float(num_with_units.replace(unit, ""))
                    return f"{round(value * factor, 2)}Gi"
                except (ValueError, TypeError):
                    return num_with_units

        return num_with_units

    @staticmethod
    def extract_jvm_params(container: Dict[str, Any]) -> str:
        """从容器配置中提取JVM参数"""
        # 检查环境变量
        env_vars = container.get("env", [])
        for env_var in env_vars:
            if env_var.get("name") in ["JAVA_OPTS", "JVM_OPTS"]:
                jvm_opts = env_var.get("value", "")
                return ResourceFormatter.extract_jvm_memory_params(jvm_opts)

        # 检查启动命令
        cmd = container.get("command", [])
        args = container.get("args", [])
        full_command = " ".join(cmd + args)
        return ResourceFormatter.extract_jvm_memory_params(full_command)

    @staticmethod
    def extract_jvm_memory_params(text: str) -> str:
        """从文本中提取JVM内存参数"""
        xms_match = re.search(r"-Xms(\d+[kKmMgGtT]?)", text)
        xmx_match = re.search(r"-Xmx(\d+[kKmMgGtT]?)", text)

        memory_parts = []
        if xms_match:
            memory_parts.append(f"Xms:{xms_match.group(1)}")
        if xmx_match:
            memory_parts.append(f"Xmx:{xmx_match.group(1)}")

        return " ".join(memory_parts) if memory_parts else ""


class ContainerFilter:
    """容器过滤器"""

    def __init__(self, config: Dict[str, Any]):
        self.config = config

    def should_include(self, container_name: str, workload_name: str, workload_kind: str, namespace: str) -> bool:
        """判断是否应该包含该容器"""
        # 检查包含规则
        for ns_pattern, wl_pattern, c_pattern in self.config.get("include_rules", []):
            if (
                self._matches_pattern(namespace, ns_pattern)
                and self._matches_pattern(workload_name, wl_pattern)
                and self._matches_pattern(container_name, c_pattern)
            ):
                return True

        # 检查排除规则
        for ns_pattern, wl_pattern, c_pattern in self.config.get("exclude_rules", []):
            if (
                self._matches_pattern(namespace, ns_pattern)
                and self._matches_pattern(workload_name, wl_pattern)
                and self._matches_pattern(container_name, c_pattern)
            ):
                return False

        # 没有排除的默认包含
        return True

    def _matches_pattern(self, value: str, pattern: str) -> bool:
        """检查值是否匹配模式"""
        return pattern == "*" or fnmatch(value, pattern)


class DataCollector:
    """数据收集器基类"""

    def __init__(self, kubectl_client: KubectlClient, logger: Logger):
        self.kubectl = kubectl_client
        self.logger = logger

    def collect(self) -> List[List[str]]:
        """收集数据，返回二维列表"""
        raise NotImplementedError


class CustomerInfoCollector(DataCollector):
    """客户信息收集器"""

    def collect(self) -> Tuple[List[List[str]], str]:
        """收集客户信息，返回数据和客户名称"""

        def valid_input(prompt: str, validator: callable, err_msg: str, default: str = None) -> str:
            while True:
                value = input(prompt).strip()
                if not value:
                    if default:
                        return default
                    else:
                        continue
                if validator(value):
                    return value
                print(err_msg)

        cust = input("请输入客户全称: ") or "匿名客户"
        env = valid_input(prompt="请输入环境类型: ", validator=lambda x: x in ["开发环境","测试环境","生产环境"], err_msg="请输入'开发环境'、'测试环境'或'生产环境'！", default = "未知环境")
        pro = valid_input(prompt="请输入专业用户数: ", validator=lambda x: x.isdigit(), err_msg="请输入有效数字！", default='未提供')
        all = valid_input(prompt="请输入全员用户数: ", validator=lambda x: x.isdigit(), err_msg="请输入有效数字！", default='未提供')
        data = [["===== 客户详情 ====="], ["客户名称", "环境类型", "专业用户", "全员用户"], [cust, env, pro, all]]

        return data, cust + '_' + env


class NodeDataCollector(DataCollector):
    """节点数据收集器"""

    def collect(self) -> List[List[str]]:
        self.logger.info("开始获取节点数据")
        nodes_data = self.kubectl.get_nodes()

        nodes_info = [
            ["===== 集群节点参数详情 ====="],
            ["Node", "Role", "OS", "Status", "VERSION", "CPU", "Memory"],
        ]

        for item in nodes_data["items"]:
            node_name = item["metadata"]["name"]

            # 节点角色
            labels = item["metadata"].get("labels", {})
            node_role = "master" if any(label in labels for label in ["node-role.kubernetes.io/control-plane", "node-role.kubernetes.io/master"]) else "worker"

            # 资源信息
            capacity = item["status"]["capacity"]
            node_cpu = capacity.get("cpu", "")
            node_mem = ResourceFormatter.convert_to_gi(capacity.get("memory", ""))

            node_info = item["status"]["nodeInfo"]
            node_os_ver = node_info["osImage"]
            node_k8s_ver = node_info["kubeletVersion"]

            # 节点状态
            node_status = "N/A"
            for condition in item["status"].get("conditions", []):
                if condition["type"] == "Ready":
                    node_status = "Ready" if condition["status"] == "True" else "NotReady"
                    break

            nodes_info.append([node_name, node_role, node_os_ver, node_status, node_k8s_ver, node_cpu, node_mem])
            self.logger.debug(f"处理节点: {node_name}")

        self.logger.info("节点数据获取完成")
        return nodes_info


class WorkloadDataCollector(DataCollector):
    """工作负载数据收集器"""

    def __init__(self, kubectl_client: KubectlClient, logger: Logger, container_filter: ContainerFilter, cnpg_config: Dict[str, Any]):
        super().__init__(kubectl_client, logger)
        self.container_filter = container_filter
        self.cnpg_config = cnpg_config

    PRIORITY_NAMESPACES = [
        "kd-cosmic-mc",
        "kd-cosmic-xk",
        "kd-cosmic-qing",
        "kd-cosmic-log",
        "kd-cosmic-ops",
        "zk-system",
        "redis-system",
        "rabbitmq-system",
        "pg-system",
        "versitygw-system",
    ]

    def collect(self) -> List[List[str]]:
        self.logger.info("开始获取工作负载数据")

        workload_lines = [
            ["===== 组件配置参数详情 ====="],
            [
                "Namespace",
                "Kind",
                "Name",
                "Replicas",
                "Container",
                "CPU_Limit",
                "MEM_Limit",
                "CPU_Request",
                "MEM_Request",
                "JVM_OPTS",
            ],
        ]

        # 处理Deployment
        deployments = self.kubectl.get_deployments()
        workload_lines.extend(self._process_workloads(deployments["items"], "Deployment"))

        # 处理StatefulSet
        statefulsets = self.kubectl.get_statefulsets()
        workload_lines.extend(self._process_workloads(statefulsets["items"], "StatefulSet"))

        # 处理CNPG Cluster
        cnpg_data = self.kubectl.get_cnpg_cluster(self.cnpg_config["namespace"], self.cnpg_config["name"])
        if cnpg_data.get("metadata"):
            workload_lines.extend(self._process_cnpg_cluster(cnpg_data))

        def sort_key(row):
            ns = row[0]
            name = row[2]
            # 联合排序键：
            # 1. 命名空间优先级（列表内按索引，列表外设为inf）
            # 2. 命名空间字母序
            # 3. workload名称字母序
            try:
                ns_priority = self.PRIORITY_NAMESPACES.index(ns)
            except ValueError:
                ns_priority = float("inf")

            return (ns_priority, ns, name)

        # 排序时保持标题行不变
        workload_lines = [workload_lines[0], workload_lines[1]] + sorted(workload_lines[2:], key=sort_key)

        self.logger.info(f"工作负载数据获取完成")
        return workload_lines

    def _process_workloads(self, items: List[Dict[str, Any]], kind: str) -> List[List[str]]:
        """处理工作负载项"""
        lines = []
        for item in items:
            ns = item["metadata"]["namespace"]
            name = item["metadata"]["name"]
            self.logger.debug(f"处理工作负载: {ns}/{name} ({kind})")
            replicas = f'{item["status"].get("readyReplicas", 0)}/{item["status"].get("replicas", 1)}'

            containers = item["spec"]["template"]["spec"].get("containers", [])
            container_data = self._process_containers(containers, name, kind, ns)

            for c in container_data:
                lines.append([ns, kind, name, replicas, c["name"], c["cpu_limit"], c["mem_limit"], c["cpu_request"], c["mem_request"], c["jvm_params"]])

        return lines

    def _process_containers(self, containers: List[Dict[str, Any]], workload_name: str, workload_kind: str, namespace: str) -> List[Dict[str, Any]]:
        """处理容器信息"""
        container_data = []
        for container in containers:
            container_name = container.get("name", "")

            if not self.container_filter.should_include(container_name, workload_name, workload_kind, namespace):
                continue

            res = container.get("resources", {})
            limits = res.get("limits", {})
            requests = res.get("requests", {})

            container_data.append(
                {
                    "name": container_name,
                    "cpu_limit": limits.get("cpu", ""),
                    "mem_limit": limits.get("memory", ""),
                    "cpu_request": requests.get("cpu", ""),
                    "mem_request": requests.get("memory", ""),
                    "jvm_params": ResourceFormatter.extract_jvm_params(container),
                }
            )

        return container_data

    def _process_cnpg_cluster(self, cnpg_data: Dict[str, Any]) -> List[List[str]]:
        """处理CNPG集群"""
        ns = cnpg_data["metadata"]["namespace"]
        name = cnpg_data["metadata"]["name"]
        self.logger.debug(f"处理CNPG集群: {ns}/{name}")
        replicas = f'{cnpg_data["status"].get("readyInstances", 1)}/{cnpg_data["status"].get("instances", 1)}'

        res = cnpg_data["spec"].get("resources", {})
        limits = res.get("limits", {})
        requests = res.get("requests", {})

        return [
            [
                ns,
                "Cluster",
                name,
                replicas,
                "postgres",
                limits.get("cpu", ""),
                limits.get("memory", ""),
                requests.get("cpu", ""),
                requests.get("memory", ""),
                "",  # CNPG没有JVM参数
            ]
        ]


class CNPGParamCollector(DataCollector):
    """CNPG参数收集器"""

    def __init__(self, kubectl_client: KubectlClient, logger: Logger, whitelist: List[str], cnpg_config: Dict[str, Any]):
        super().__init__(kubectl_client, logger)
        self.whitelist = whitelist
        self.cnpg_config = cnpg_config

    def collect(self) -> List[List[str]]:
        self.logger.info("开始处理CNPG参数")

        cnpg_lines = [["===== CNPG集群参数情况 ====="]]

        cnpg_data = self.kubectl.get_cnpg_cluster(self.cnpg_config["namespace"], self.cnpg_config["name"])

        if cnpg_data.get("spec"):
            try:
                params = cnpg_data["spec"]["postgresql"]["parameters"]
                filtered_params = {k: v for k, v in params.items() if k in self.whitelist}
                # cnpg_lines.append(["Param", "Value"])
                # 按whitelist顺序排序
                sorted_params = {k: filtered_params[k] for k in self.whitelist if k in filtered_params}
                cnpg_lines.extend([key, value] for key, value in sorted_params.items())
            except KeyError as e:
                self.logger.error(f"CNPG参数格式错误: {e}")
                cnpg_lines.append(["参数解析失败"])
        else:
            cnpg_lines.append(["未找到CNPG Cluster"])

        self.logger.info("CNPG参数处理完成")
        return cnpg_lines


class Exporter:
    """文件导出器"""

    def __init__(self, output_dir: str, logger: Logger):
        self.output_dir = output_dir
        self.logger = logger

    def export(self, data: Dict[str, List[List[str]]], customer_name: str) -> str:
        # 处理输出目录路径
        output_dir = os.path.abspath(self.output_dir) if self.output_dir else os.getcwd()

        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)

        # 生成文件名
        clean_name = re.sub(r"[^\w\u4e00-\u9fff]", "_", customer_name).strip("_")
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{clean_name}_cluster_info_{timestamp}.csv"
        output_path = os.path.join(output_dir, filename)

        try:
            with open(output_path, "w", newline="", encoding="utf-8-sig") as f:
                writer = csv.writer(f, delimiter=",")

                # 按顺序写入各部分数据
                sections = ["customer_info", "node_data", "workload_data", "cnpg_params"]
                for section in sections:
                    if section in data:
                        for line in data[section]:
                            writer.writerow(line)
                        writer.writerow([])  # 添加空行分隔

            # 显示完整输出路径
            abs_path = os.path.abspath(output_path)
            self.logger.info(f"数据已写入文件: {abs_path}")
            return abs_path
        except Exception as e:
            self.logger.error(f"写入文件失败: {e}")
            raise


def main():
    """主函数"""
    # 初始化组件
    logger = Logger(debug_mode=CONFIG["debug"])
    kubectl_client = KubectlClient(logger)
    container_filter = ContainerFilter(CONFIG["container_filters"])
    exporter = Exporter("", logger)

    # 收集客户信息（返回数据和客户名称）
    customer_collector = CustomerInfoCollector(kubectl_client, logger)
    customer_data, customer_name = customer_collector.collect()

    # 收集其他数据
    node_collector = NodeDataCollector(kubectl_client, logger)
    workload_collector = WorkloadDataCollector(kubectl_client, logger, container_filter, CONFIG["cnpg_cluster"])
    cnpg_param_collector = CNPGParamCollector(kubectl_client, logger, CONFIG["cnpg_whitelist"], CONFIG["cnpg_cluster"])

    collected_data = {
        "customer_info": customer_data,
        "node_data": node_collector.collect(),
        "workload_data": workload_collector.collect(),
        "cnpg_params": cnpg_param_collector.collect(),
    }

    # 导出数据
    exporter.export(collected_data, customer_name)


if __name__ == "__main__":
    main()
